#!/bin/env python

import matplotlib
matplotlib.use('Agg')
from plot_profiles_functions import *
from matplotlib import rc
from matplotlib.font_manager import FontProperties
matplotlib.rc('xtick',labelsize='x-large')
matplotlib.rc('ytick',labelsize='x-large')
# matplotlib.rc('axes.formatter',limits='-2,3')

# from pylab import *
import sys, os
import pylab as p
import numpy as np

sys.path.append('/home/ca347/codes/python/')

## This takes in a text file of run directories, aexpns, legendlabels, and linestyles - makes multiple plots ##

############# Change here #############

fT = {'saveas':'fT', 'rvar':0, 'gasvar':1, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_a', 'num_bins':99,'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{T}\ [\mathrm{K}]$','xlog':1, 'ylog':1, 'xlims':(5e-3,5.0),'ylims':(1e5,1e9),'legend_position':'best','title':'Temperature'}
fP = {'saveas':'fP', 'rvar':0, 'gasvar':2, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{P}\ [\mathrm{ergs}/\mathrm{cm}^3]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-15,1e-8),'legend_position':'best','title':'Pressure'}
fK = {'saveas': 'fK', 'rvar':0, 'gasvar':3, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{K}\ [\mathrm{keV}\ \mathrm{cm}^2]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-1,1e5),'legend_position':'best','title':'Entropy'}
fTmass = {'saveas':'fTmass', 'rvar':0, 'gasvar':1, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_massweighted_a', 'num_bins':99,'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{T}\ [\mathrm{K}]$','xlog':1, 'ylog':1, 'xlims':(5e-3,5.0),'ylims':(1e5,1e9),'legend_position':'best','title':'Temperature (Mass weighted)'}
fPmass = {'saveas':'fPmass', 'rvar':0, 'gasvar':2, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_massweighted_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{P}\ [\mathrm{ergs}/\mathrm{cm}^3]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-15,1e-8),'legend_position':'best','title':'Pressure (Mass weighted)'}
fKmass = {'saveas': 'fKmass', 'rvar':0, 'gasvar':3, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_massweighted_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{K}\ [\mathrm{keV}\ \mathrm{cm}^2]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-1,1e5),'legend_position':'best','title':'Entropy (Mass weighted)'}

fTnorm500 = {'saveas':'fTnorm500', 'rvar':0, 'gasvar':1, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_a', 'num_bins':99,'rbins':arange(0.0,3.,0.01), # 10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{T}/\mathrm{T}_{500}$','xlog':0, 'ylog':0, 'xlims':(5e-3,1.5),'ylims':(1e-1,3.0),'legend_position':'best','title':'Temperature'}
fPnorm500 = {'saveas':'fPnorm500', 'rvar':0, 'gasvar':2, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{P}/\mathrm{P}_{500}$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-2,1e4),'legend_position':'best','title':'Pressure'}
fKnorm500 = {'saveas': 'fKnorm500', 'rvar':0, 'gasvar':3, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{K}/\mathrm{K}_{500}$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(5e-3,1e2),'legend_position':'best','title':'Entropy'}
fTmassnorm500 = {'saveas':'fTmassnorm500', 'rvar':0, 'gasvar':1, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_massweighted_a', 'num_bins':99,'rbins':arange(0.0,3.0,0.01), # 10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{T}/\mathrm{T}_{500}$','xlog':0, 'ylog':0, 'xlims':(5e-3,1.5),'ylims':(1e-1,3.0),'legend_position':'best','title':'Temperature (Mass weighted)'}
fPmassnorm500 = {'saveas':'fPmassnorm500', 'rvar':0, 'gasvar':2, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_massweighted_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{P}/\mathrm{P}_{500}$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-3,1e4),'legend_position':'best','title':'Pressure (Mass weighted)'}
fKmassnorm500 = {'saveas': 'fKmassnorm500', 'rvar':0, 'gasvar':3, 'load':'icm', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bgpro_massweighted_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{K}/\mathrm{K}_{500}$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(5e-3,1e2),'legend_position':'best','title':'Entropy (Mass weighted)'}

# Density profiles
frhoDM = {'saveas': 'frhoDM', 'rvar':0, 'gasvar':1, 'load':'density', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\\rho_{\mathrm{DM}}\ [\mathrm{M}_{\odot}/\mathrm{h}(\mathrm{kpc}/\mathrm{h})^{-3}]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e1,1e10),'legend_position':'best','title':'DM density'}
fr2rhoDM = {'saveas': 'fr2rhoDM', 'rvar':0, 'gasvar':1, 'load':'density', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{r}^2\\rho_{\mathrm{DM}}\ [\mathrm{kpc}^2\mathrm{M}_{\odot}/\mathrm{h}(\mathrm{kpc}/\mathrm{h})^{-3}]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e8,1e11),'legend_position':'best','title':'DM density'}
frhog = {'saveas': 'frhog', 'rvar':0, 'gasvar':2, 'load':'density', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\\rho_{\mathrm{icm}}\ [\mathrm{M}_{\odot}/\mathrm{h}(\mathrm{kpc}/\mathrm{h})^{-3}]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e3,1e10),'legend_position':'best','title':'Gas density'}
fnH = {'saveas': 'fnH', 'rvar':0, 'gasvar':2, 'load':'density', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{n}_{\mathrm{H}}\mathrm{E}^{-2}(z)\ [\mathrm{cm}^{-3}]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-5,1e2),'legend_position':'best','title':'Hydrogen number density'}

# Enclosed masses
fMgas = {'saveas': 'fMgas', 'rvar':0, 'gasvar':2, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{M}_{\mathrm{gas}}(<R)\ [\mathrm{M}_{\odot}/h]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e10,1e14),'legend_position':'best','title':'Enclosed Gas Mass'}
fMstar = {'saveas': 'fMstar', 'rvar':0, 'gasvar':4, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{M}_{\mathrm{*}}(<R)\ [\mathrm{M}_{\odot}/h]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e11,1e14),'legend_position':'best','title':'Enclosed Stellar Mass'}
fMcg = {'saveas': 'fMcg', 'rvar':0, 'gasvar':5, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{M}_{\mathrm{cg}}(<R)\ [\mathrm{M}_{\odot}/h]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e10,1e14),'legend_position':'best','title':'Enclosed Cold Gas Mass'}

# mass fractions
fgasfraction = {'saveas': 'fgasfraction', 'rvar':0, 'gasvar':2, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{g}}(<R) [\\Omega_b/\\Omega_m]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-1,5.0),'legend_position':'best','title':'Gas Fraction'}
fstarfraction = {'saveas': 'fstarfraction', 'rvar':0, 'gasvar':4, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{*}}(<R) [\\Omega_b/\\Omega_m]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-1,10),'legend_position':'best','title':'Stellar Fraction'}
fcgfraction = {'saveas': 'fcgfraction', 'rvar':0, 'gasvar':5, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{cg}}(<R) [\\Omega_b/\\Omega_m]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-2,5.0),'legend_position':'best','title':'Cold Gas Fraction'}
fcgstarfraction = {'saveas': 'fcgstarfraction', 'rvar':0, 'gasvar':4, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{cg}+\mathrm{*}}(<R) [\\Omega_b/\\Omega_m]$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-1,10),'legend_position':'best','title':'Cold Gas and Star Fraction'}

# Mass fraction of baryons
fgasbary = {'saveas': 'fgasbary', 'rvar':0, 'gasvar':2, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{g}}(<R)/\mathrm{f}_{\mathrm{b}}(<R)$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-2,1.0),'legend_position':'best','title':'Gas Fraction of Baryons'}
fstarbary = {'saveas': 'fstarbary', 'rvar':0, 'gasvar':4, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{*}}(<R)/\mathrm{f}_{\mathrm{b}}(<R)$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-1,1.0),'legend_position':'best','title':'Stellar Fraction of Baryons'}
fcgbary = {'saveas': 'fcgbary', 'rvar':0, 'gasvar':5, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{cg}}(<R)/\mathrm{f}_{\mathrm{b}}(<R)$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-2,1.0),'legend_position':'best','title':'Cold Gas Fraction of Baryons'}
fcgstarbary = {'saveas': 'fcgstarbary', 'rvar':0, 'gasvar':4, 'load':'cumulative', 'profile_loc':'/profiles_hlist','profile_prefix':'/h_bmpro_a', 'num_bins':99, 'rbins':10.**arange(-3,0.7,0.01),
      'xlabel':'$\mathrm{R}/\mathrm{R}_{500}$','ylabel':'$\mathrm{f}_{\mathrm{cg}+\mathrm{*}}(<R)/\mathrm{f}_{\mathrm{b}}(<R)$','xlog':1,'ylog':1,'xlims':(5e-3,5.0),'ylims':(1e-2,1.0),'legend_position':'best','title':'Cold Gas and Star Fraction of Baryons'}

#################CHECKED TO HERE##############
fTnormTmg = {'saveas':'fTnormTmg', 'gasvar':1, 'yname':'$\mathrm{T}/T_{\mathrm{mg}}$','ylims':(0.3,3),'title':'Temperature'}
frhognormrhoc = {'saveas': 'frhognormrhoc', 'gasvar':2, 'yname':'$\\rho_{g}/\\rho_{c}$','ylims':(1.,1e6),'title':'Gas Density'}
fgasdiffr = {'saveas': 'fgasdiffr', 'gasvar':2, 'yname':'$f_{\mathrm{bar,gas}}(r)\ [\Omega_b/\Omega_m]$','ylims':(),'title':'Differential Gas Fraction'}
fstareff = {'saveas': 'fstareff', 'gasvar':2, 'yname':'$\mathrm{M}_{\mathrm{*}}/(\mathrm{M}_{\mathrm{cg}}+\mathrm{M}_{\mathrm{*}})$','ylims':(0.84,1.05),'title':'Stellar Mass Efficiency'}
ftcool = {'saveas': 'ftcool', 'gasvar':4, 'yname':'tcool [yr]','ylims':(1e3,1e12),'title':'Cooling Time'}
fxrayem = {'saveas': 'fxrayem', 'gasvar':2, 'yname':'$n_{\mathrm{H}}^2*(T/T_{500})^{1/2}(r) [cm^{-3}]$','ylims':(),'title':'Differential X-ray Flux'}

# Below is the list of items we wish to plot from the above dictionaries
#plot_items = [ fMstar, fstarfraction, fstarbary ] # Can't run this on NR profiles
# plot_items = [ fT, fP, fK, fTmass, fPmass, fKmass, fTnorm500, fPnorm500, fKnorm500, fTmassnorm500, fPmassnorm500, fKmassnorm500,
# 	       frhoDM, fr2rhoDM, frhog, fnH, fMgas, fMcg,
# 	       fgasfraction, fcgfraction, fgasbary, fcgbary ] # Can run on NR profiles
plot_items = [ fT, fP, fK, fTmass, fPmass, fKmass, fTnorm500, fPnorm500, fKnorm500, fTmassnorm500, fPmassnorm500, fKmassnorm500,
  	       frhoDM, fr2rhoDM, frhog, fnH, fMgas, fMstar, fMcg,
 	       fgasfraction, fstarfraction, fcgfraction, fcgstarfraction, fgasbary, fstarbary, fcgbary, fcgstarbary ]
# plot_items = [ fK, fKmass, fKnorm500, fKmassnorm500, fP, fPmass, fPnorm500, fPmassnorm500, frhoDM, frhog, fnH ]
#######################################

class MyLogFormatter(LogFormatter) :

	def __call__(self,x,pos=None) :
		if abs(log10(x)) <2 :
			return "$%g$" % (x,)
		else :
			return "$10^{%g}$" % (log10(x),)

if len(sys.argv) != 2 :
	print 'Usage: plot.py listofdirectories.txt[rundir, aexpns, legendlabel,markercolor/type and halo_id]'
	sys.exit(1)

rundirlist = sys.argv[1]

# Parse the input file for rundirectory locations, aexpns, marker style/color, and halo_ids - these are dictionaries labeled with legend_labels for each data set
rundirs, aexpns, markers, halo_ids, png_folder = parse_plot_input( rundirlist )
print "rundirs:", rundirs
print "aexpns:", aexpns
print "markers:", markers
print "halo_ids:", halo_ids

### PLOT EACH PLOT IN LIST OF PLOT ITEMS ###
		   
for plot_item in plot_items :
	# Create a 7x7 figure
	f = p.figure( figsize = (7,7) )
	axes([0.20,0.15,0.75,0.75])
	f.gca().axvline(x=1.0,linestyle='dotted',linewidth=1.5,color='k')
	plots = {}
	# Loop over each data set (uniquely separated by legend_labels, which are the keys the dictionaries)
	for legend_label in rundirs.keys() :
		print "Plotting for :" + legend_label
		print rundirs[legend_label]
		print 'aexp:'+aexpns[legend_label]+', halo_id:',halo_ids[legend_label]
		interpolated_profile = plot_out( rundirs[legend_label],
					 aexpns[legend_label],
					 halo_ids[legend_label],
					 plot_item )  
		plots[legend_label] = f.gca().plot( plot_item['rbins'], interpolated_profile, markers[legend_label], linewidth=0.5 ) # Linewidth normaly 2.0

	# Specify plot settings
	f.gca().set_xlabel(plot_item['xlabel'],fontsize='xx-large')
	f.gca().set_ylabel(plot_item['ylabel'],fontsize='xx-large')
	if (plot_item['xlog']) :
		f.gca().set_xscale('log')
		f.gca().xaxis.set_major_formatter(MyLogFormatter())
	if (plot_item['ylog']) :
		f.gca().set_yscale('log')
		f.gca().yaxis.set_major_formatter(MyLogFormatter())
	if plot_item['xlims'] != 'None' :
		f.gca().set_xlim(plot_item['xlims'])
	if plot_item['ylims'] != 'None' :
		f.gca().set_ylim(plot_item['ylims'])

	if len(plots.keys()[0]) > 2 :
		f.gca().legend([plots[legend_label] for legend_label in plots.keys()],plots.keys(),plot_item['legend_position'],markerscale=7.)
		leg = gca().get_legend()
		legtext = leg.get_texts()
		leglines = leg.get_lines()
		setp(legtext,fontsize='large')
		setp(leglines,linewidth=2.0)
		leg.draw_frame(False)

	for tick in f.gca().yaxis.get_major_ticks() :
		tick.label1.set_fontsize('x-large')
	for tick in f.gca().xaxis.get_major_ticks() :
		tick.label1.set_fontsize('x-large')
	for line in f.gca().xaxis.get_ticklines()+f.gca().yaxis.get_ticklines() :
		line.set_markersize(10)
		line.set_markeredgewidth(1.2)
	
	f.gca().set_title(plot_item['title'],fontsize='xx-large')
	f.gca().xaxis.set_major_formatter(MyLogFormatter())
	f.gca().yaxis.set_major_formatter(MyLogFormatter())
	
	# Save the figure
	f.savefig( png_folder+'/'+plot_item['saveas'], dpi=100 )








